package ml.features;

import java.io.BufferedWriter;
import java.io.File;
import java.io.FileWriter;
import java.io.IOException;
import java.util.ArrayList;
import java.util.Arrays;
import java.util.List;
import java.util.Scanner;
import java.util.TreeSet;

import corpus.SentPolarityCorpus;

import ml.Data;
import parser.LTHSRLAnnotation;
import parser.ParserAnnotation;
import parser.ParserPipeline;

public class SRLExtractor extends FeatureExtractor {
	private String[] srls;
	private String[] deprels;
	private int maxHead = 1;

	public SRLExtractor(int offset) throws IOException {
		super(offset);
		srls = loadList(new File("features/srl"));
		deprels = loadList(new File("features/deprel"));
		Scanner scr = new Scanner(new File("features/headersize"));
		maxHead = scr.nextInt();
		scr.close();
	}

	public SRLExtractor(int offset, SentPolarityCorpus corpus, ParserPipeline p) {
		super(offset);
		TreeSet<String> srlLabels = new TreeSet<String>();
		TreeSet<String> deprelLabels = new TreeSet<String>();
		for (String sent : corpus.getCorpus().keySet()) {
			LTHSRLAnnotation pa = (LTHSRLAnnotation) (corpus.getParsedPA(sent,
					p));
			srlLabels.addAll(pa.getSRL(0));
			deprelLabels.addAll(pa.getDeprel(0));
			for (int i : pa.getHeadIds(0)) {
				if (i > maxHead)
					maxHead = i;
			}
		}
		srls = srlLabels.toArray(new String[srlLabels.size()]);
		deprels = deprelLabels.toArray(new String[deprelLabels.size()]);
	}

	@Override
	public List<Data> extract(ParserAnnotation pa, int sentence) {
		throw new Error("SRL features can only be used by sequence methods.");
	}

	@Override
	public List<Data>[] extractSequence(ParserAnnotation pa, int sentence) {
		LTHSRLAnnotation la = (LTHSRLAnnotation) pa;
		int tokens = pa.getTokens(sentence).size();
		List<Data>[] datas = new List[1 + srls.length + deprels.length];
		for (int i = 0; i < datas.length; i++)
			datas[i] = new ArrayList<Data>();

		for (int i = 0; i < tokens; i++) {
			double value = la.getHeadIds(sentence).get(i);
			value = normalize(0, maxHead, value);
			datas[0].add(new Data(value, i + offset));
		}

		for (String deprel : la.getDeprel(0)) {
			for (int i = 0; i < deprels.length; i++) {
				String sdeprel = deprels[i];
				datas[1 + i].add(new Data(sdeprel.equals(deprel), 1 + i
						+ offset));
			}
		}
		
		for (String srl : la.getSRL(0)) {
			for (int i = 0; i < srls.length; i++) {
				String ssrl = srls[i];
				datas[1 + deprels.length + i].add(new Data(ssrl.equals(srl), 1
						+ deprels.length + i + offset));
			}
		}
		return datas;
	}

	@Override
	public int nbrColumns() {
		return 1 + srls.length + deprels.length;
	}

	@Override
	public String[] wekaDescription() {
		throw new Error("SRL features can only be used by sequence methods.");
	}

	@Override
	public void store() throws IOException {
		File dfile = new File("features/deprel");
		File sfile = new File("features/srl");
		storeSet(Arrays.asList(deprels), dfile);
		storeSet(Arrays.asList(srls), sfile);
		File hfile = new File("features/headersize");
		BufferedWriter bw = new BufferedWriter(new FileWriter(hfile));
		bw.append("" + maxHead);
		bw.flush();
		bw.close();
	}
}
